#include <assert.h>
#include <iostream>
#include <WinSock2.h>
#include <Ws2tcpip.h>

#include "SocketServer.h"

//===========================================================================
// Socket
//===========================================================================

//===========================================================================
Socket::Socket( void )
    : m_socket( INVALID_SOCKET )
{}

//===========================================================================
Socket::Socket(
    const Socket & socket
) {
    this->m_socket = socket.GetSocket();
}

//===========================================================================
Socket::~Socket( void )
{
}

//===========================================================================
void Socket::Initialize(
    SOCKET socket
) {
    this->Release();

    this->m_socket = socket;
    this->m_buffer.Initialize();
}

//===========================================================================
SOCKET Socket::GetSocket( void ) const
{
    return this->m_socket;
}

//===========================================================================
bool Socket::Recv( void )
{
    assert( INVALID_SOCKET != this->m_socket );

    DWORD received = 0;
    DWORD flag     = 0;
    if( SOCKET_ERROR == ::WSARecv(
            this->m_socket,
            &this->m_buffer.buf,
            1,
            &received,
            &flag,
            NULL,
            NULL
    ) )
        if( WSAEWOULDBLOCK != ::WSAGetLastError() )
            return false;

    if( !this->m_buffer.Recv( received ) )
        return false;

    return true;
}

//===========================================================================
void Socket::Drain(
    unsigned size
) {
    this->m_buffer.Drain( size );
}

//===========================================================================
const char * Socket::GetBuffer( void ) {
    return this->m_buffer.buffer;
}

//===========================================================================
unsigned int Socket::GetBufferSize( void ) {
    return this->m_buffer.size;
}

//===========================================================================
void Socket::Release( void )
{
    if( INVALID_SOCKET != this->m_socket ) {
        ::closesocket( this->m_socket );
        this->m_socket = INVALID_SOCKET;
    }
}

//===========================================================================
// Clients
//===========================================================================

//===========================================================================
Clients::Clients( void )
    : m_count( 0 )
{}

//===========================================================================
Clients::~Clients( void ) {
    this->Cleanup();
}

//===========================================================================
unsigned int Clients::GetCount( void ) const {
    return this->m_count;
}

//===========================================================================
const WSAEVENT * Clients::GetEvents( void ) const {
    return this->m_events;
}

//===========================================================================
Socket * Clients::GetSocket(
    unsigned int index
) {
    return index < this->m_count
         ? &this->m_sockets[ index ]
         : NULL;
}

//===========================================================================
WSAEVENT * Clients::GetEvent(
    unsigned int index
) {
    return index < this->m_count
         ? &this->m_events[ index ]
         : NULL;
}

//===========================================================================
int Clients::Generate(
    SOCKET socket,
    long   events /* = 0 */
) {
    if( INVALID_SOCKET == socket )
        return -1;

    if( !( this->m_count + 1 < MAX_CLIENTS ) )
        return -2;
    
    WSAEVENT & event = this->m_events[ this->m_count ];
    event = ::WSACreateEvent();
    if( WSA_INVALID_EVENT == event )
        return -3;

    if( events )
        ::WSAEventSelect(
            socket,
            event,
            events
        );

    this->m_sockets[ this->m_count ].Initialize( socket );    
    this->m_count++;

    return this->m_count - 1;
}

//===========================================================================
void Clients::Release(
    unsigned int index
) {
    assert( index < this->m_count );

    this->m_sockets[ index ].Release();
    ::WSACloseEvent( this->m_events[ index ] );
    this->m_events[ index ] = WSA_INVALID_EVENT;

    if( index < this->m_count - 1 ) {
        this->m_sockets[ index ] = this->m_sockets[ this->m_count - 1 ];
        this->m_events[  index ] = this->m_events[  this->m_count - 1 ];    
        this->m_events[ this->m_count - 1 ] = WSA_INVALID_EVENT;
    }
    
    this->m_count--;
}

//===========================================================================
void Clients::Cleanup( void ) {
    if( this->m_count ) {
        unsigned index = this->m_count - 1;
        while( true ) {
            this->Release( index );
            if( !index )
                break;
            index--;
        }
    }
}

//===========================================================================
// SocketServer
//===========================================================================

//===========================================================================
SocketServer::SocketServer( void )
    : m_port(  0                     )
    , m_thread( INVALID_HANDLE_VALUE )
    , m_stop(   INVALID_HANDLE_VALUE )
{
    WSADATA wsa;
    int result = ::WSAStartup( MAKEWORD( 2, 2 ), &wsa );
    assert( !result );
}

//===========================================================================
SocketServer::~SocketServer( void )
{
    if( this->IsRunning() )
        this->Stop();
    
    ::WSACleanup();
}

//===========================================================================
bool SocketServer::Initialize(
    unsigned int port
) {
    this->m_port = port;

    return true;
}

//===========================================================================
bool SocketServer::IsRunning( void )
{
    return this->m_thread != INVALID_HANDLE_VALUE;
}

//===========================================================================
bool SocketServer::Run( void )
{
    int index = this->m_clients.Generate(
        ::socket( AF_INET, SOCK_STREAM, 0 ),
        FD_ACCEPT | FD_CLOSE
    );
    if( index < 0 )
        return false;
    SOCKET socket = this->m_clients.GetSocket( index )->GetSocket();

    SOCKADDR_IN address;
    address.sin_family      = AF_INET;
    address.sin_addr.s_addr = htonl( INADDR_ANY );
    address.sin_port        = htons( this->m_port );

    if( SOCKET_ERROR == ::bind(
            socket,
            reinterpret_cast< PSOCKADDR >( &address ),
            sizeof( address )
    ) )
        return false;
    
    if( listen( socket, 5 ) )
        return false;
    
    this->m_thread = ::CreateThread(
        NULL,
        0,
        SocketServer::EventThreadProc,
        this,
        0,
        NULL
    );

    this->m_stop = ::CreateEvent(
        NULL, TRUE, FALSE, TEXT( "ThreadStopEvent" )
    );

    return true;
}

//===========================================================================
Clients * SocketServer::GetClients( void ) {
    return &this->m_clients;
}

//===========================================================================
bool SocketServer::IsTimeToStop( void )
{
    return WSA_WAIT_EVENT_0 == ::WaitForSingleObject( this->m_stop, 20 )
         ? true
         : false;
}

//===========================================================================
void SocketServer::Stop( void )
{
    if( !this->IsRunning() )
        return;
    
    ::SetEvent( this->m_stop );

    ::WaitForSingleObject( this->m_thread, INFINITE );
    ::CloseHandle( this->m_thread );
    this->m_thread = INVALID_HANDLE_VALUE;

    ::CloseHandle( this->m_stop );
    this->m_stop = INVALID_HANDLE_VALUE;
    
}

//===========================================================================
DWORD WINAPI SocketServer::EventThreadProc(
    LPVOID param
) {
    SocketServer * server = static_cast< SocketServer * >( param );
    if( !server )
        return -1;

    Clients * clients = server->GetClients();
    while( true ) {
        if( server->IsTimeToStop() )
            return 0;

        const WSAEVENT * events = clients->GetEvents();
              DWORD      count  = clients->GetCount();
        DWORD index = ::WSAWaitForMultipleEvents(
            count,
            events,
            false,
            20,
            false
        );
        if( WSA_WAIT_FAILED == index )
            return -2;

        if( WSA_WAIT_TIMEOUT == index )
            continue;

        WSANETWORKEVENTS event;
        index -= WSA_WAIT_EVENT_0;
        Socket *   socket = clients->GetSocket( index );
        WSAEVENT * handle = clients->GetEvent(  index );
        if( SOCKET_ERROR == ::WSAEnumNetworkEvents(
                socket->GetSocket(),
                *handle,
                &event
        ) )
            return -3;

        // accept
        if( FD_ACCEPT & event.lNetworkEvents )
            if( !event.iErrorCode[ FD_ACCEPT_BIT ] ) {
                int accepted = clients->Generate(
                    ::accept( socket->GetSocket(), NULL, NULL ),
                    FD_READ | FD_WRITE | FD_CLOSE
                );
                if( accepted > -1 )
                    server->OnConnected(
                        clients->GetSocket( accepted )->GetSocket()
                    );
            }

        // read
        if( FD_READ & event.lNetworkEvents )
            if( !event.iErrorCode[ FD_READ_BIT ] )
                if( socket->Recv() ) {
                    unsigned processed = 0;
                    server->OnRead(
                        socket->GetSocket(),
                        socket->GetBuffer(),
                        socket->GetBufferSize(),
                        processed
                    );
                    if( processed > 0 )
                        socket->Drain( processed );
                }

        // write
        if( FD_WRITE & event.lNetworkEvents ) {

        }

        // close
        if( FD_CLOSE & event.lNetworkEvents ) {
            server->OnDisconnected( socket->GetSocket() );
            clients->Release( index );
        }
    }
    
    return 0;
}
